[Perf][Spec Decode] Fuse target temperature in rejection sampler - #53090

Summary

Fuse target-temperature scaling into the MRV2 rejection-sampling kernels for temperature-only requests. This avoids materializing an FP32 target-logits buffer and removes chunking on this path.

Requests using other logits processors or processed logprobs keep the existing fallback. Startup warmup compiles the fused variants to avoid first-request JIT compilation.

Results

L40S, BF16 target logits, vocabulary 151936, speculative K=4:

At batch 128, peak temporary allocation decreased from 371.4 MiB to 0.5 MiB (370.9 MiB saved).

Two five-round Qwen3.5-9B end-to-end comparisons (ngram K=4, 128 prompts, 64 output tokens) remained within run-to-run noise: -0.66% and +0.99%. These are rejection-sampler gains; no end-to-end throughput improvement is claimed.

A fresh Triton-cache run confirmed that the fused rejection-sampler kernels do not compile during inference after startup warmup.

Tests

.venv/bin/python -m pytest -q \
  tests/v1/spec_decode/test_rejection_sampler_utils.py \
  tests/v1/worker/test_gpu_rejection_sampler_chunking.py \
  tests/v1/worker/test_gpu_sampler_flags.py \
  tests/v1/worker/test_gpu_rejection_sampler_i64.py \
  tests/v1/worker/test_gpu_warmup_blocks.py
# 125 passed

uv run --no-sync ruff check $(git diff --name-only origin/main...HEAD -- '*.py')
uv run --no-sync ruff format --check $(git diff --name-only origin/main...HEAD -- '*.py')

Coverage includes FP32/FP16/BF16 target logits, one-hot/full draft logits, standard/block/synthetic verification, mixed temperatures, multi-block vocabularies, and fallback paths.

Benchmark commands:

.venv/bin/python benchmarks/kernels/benchmark_rejection_sampler_temperature.py \
  --batch-sizes 1 8 32 128 --warmup-ms 100 --rep-ms 300
.venv/bin/python benchmarks/kernels/benchmark_rejection_sampler_temperature.py \
  --batch-sizes 1 8 32 128 --warmup-ms 100 --rep-ms 300 --with-draft-logits

Duplicate check

No equivalent open PR was found. #41258 and #45369 target the legacy sampler. #48928 targets FlashInfer AIR top-p sampling, while this fast path excludes top-p.

AI assistance

AI assistance was used for implementation and validation. The submitter reviewed every changed line and ran the tests above.